package cn.codecrab.ai.langchain4j.config;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.elasticsearch.ElasticsearchEmbeddingStore;
import dev.langchain4j.store.embedding.pinecone.PineconeEmbeddingStore;
import dev.langchain4j.store.embedding.pinecone.PineconeServerlessIndexConfig;
import org.elasticsearch.client.RestClient;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * @author 王刚
 * @since 2025年05月11日 18:18
 */
@Configuration
public class EmbeddingStoreConfig {

    @Bean
    public EmbeddingStore<TextSegment> esEmbeddingStore(RestClient restClient) {
        return ElasticsearchEmbeddingStore.builder().restClient(restClient).indexName("xiaozhi-index").build();
    }

    /**
     * 构建这个Bean的时候，会去pinecone创建一个index指定的索引
     */
    @Bean
    public EmbeddingStore<TextSegment> pineconeEmbeddingStore(EmbeddingModel embeddingModel) {
        return PineconeEmbeddingStore.builder()
            .apiKey(System.getenv("PINECONE_API_KEY"))
            .index("xiaozhi-index") // 如果指定的索引不存在，将创建一个新的索引
            .nameSpace("xiaozhi-namespace") // 指定的名称空间不存在，将创建一个新的名称空间
            .createIndex(PineconeServerlessIndexConfig.builder()
                .dimension(embeddingModel.dimension()) // 指定索引的向量维度，该维度与embeddedModel生成的向量维度相同。
                .cloud("AWS") // 指定索引部署在AWS云服务上。
                .region("us-east-1") // 指定索引所在的 AWS 区域为 us-east-1。
                .build()
            )
            .build();
    }

}
